{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1e076089-0675-46e3-9702-dfe256edeb32",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import os\n",
    "import sys\n",
    "import random\n",
    "from dataclasses import dataclass, field\n",
    "from typing import Optional\n",
    "\n",
    "import datasets\n",
    "import torch\n",
    "import numpy as np\n",
    "from datasets import ClassLabel, load_dataset, load_metric\n",
    "from datasets import load_from_disk\n",
    "\n",
    "import transformers\n",
    "from transformers import (\n",
    "    AutoConfig,\n",
    "    AutoModelForTokenClassification,\n",
    "    AutoTokenizer,\n",
    "    DataCollatorForTokenClassification,\n",
    "    HfArgumentParser,\n",
    "    PreTrainedTokenizerFast,\n",
    "    Trainer,\n",
    "    TrainingArguments,\n",
    "    set_seed,\n",
    ")\n",
    "from transformers.trainer_utils import get_last_checkpoint\n",
    "from transformers.utils import check_min_version\n",
    "from transformers.utils.versions import require_version\n",
    "\n",
    "# Model\n",
    "model_name_or_path = \"../pretrained_models/bert-base-chinese/\"\n",
    "config_name = None\n",
    "tokenizer_name = None\n",
    "cache_dir = None\n",
    "model_revision = \"main\"\n",
    "use_auth_token = False\n",
    "\n",
    "# Data\n",
    "batch_size = 8\n",
    "gradient_accumulation_steps = 64 // batch_size\n",
    "\n",
    "task_name = \"ner\"\n",
    "dataset_name = \"../data/char_ner_dataset.py\"\n",
    "dataset_config_name = \"train\" # \"test\"\n",
    "preprocessing_num_workers = 4\n",
    "pad_to_max_length = False\n",
    "max_train_samples = None\n",
    "label_all_tokens = True\n",
    "return_entity_level_metrics = True\n",
    "\n",
    "set_seed(402)\n",
    "\n",
    "names = ['B-GPE',\n",
    " 'B-LOC',\n",
    " 'B-ORG',\n",
    " 'B-PER',\n",
    " 'E-GPE',\n",
    " 'E-LOC',\n",
    " 'E-ORG',\n",
    " 'E-PER',\n",
    " 'M-GPE',\n",
    " 'M-LOC',\n",
    " 'M-ORG',\n",
    " 'M-PER',\n",
    " 'O',\n",
    " 'S-GPE',\n",
    " 'S-LOC',\n",
    " 'S-ORG',\n",
    " 'S-PER']\n",
    "\n",
    "\n",
    "# 数据\n",
    "def get_dataset(seed, test_size=0.2):\n",
    "    raw_dataset = load_dataset(dataset_name, name=dataset_config_name)\n",
    "    raw_dataset = raw_dataset[\"train\"].train_test_split(test_size=0.2, seed=seed)  # 402， 1218， 1007\n",
    "\n",
    "    return raw_dataset\n",
    "\n",
    "# 模型\n",
    "def get_model(seed):\n",
    "    config = AutoConfig.from_pretrained(model_name_or_path)\n",
    "    model = AutoModelForTokenClassification.from_pretrained(model_name_or_path, config=config)\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, config=config)\n",
    "    model.resize_token_embeddings(len(tokenizer))\n",
    "\n",
    "    raw_dataset = get_dataset(seed=seed)\n",
    "    def tokenize_and_align_labels(examples):\n",
    "        tokenized_inputs = tokenizer(examples[\"tokens\"], is_split_into_words=True,\n",
    "                                     padding=True, truncation=True)\n",
    "\n",
    "        labels = []\n",
    "        for i, label in enumerate(examples[\"labels\"]):\n",
    "            word_ids = tokenized_inputs.word_ids(batch_index=i)\n",
    "            previous_word_idx = None\n",
    "            label_ids = []\n",
    "            for word_idx in word_ids:\n",
    "                # Special tokens have a word id that is None. We set the label to -100 so they are automatically\n",
    "                # ignored in the loss function.\n",
    "                if word_idx is None:\n",
    "                    label_ids.append(-100)\n",
    "                # We set the label for the first token of each word.\n",
    "                elif word_idx != previous_word_idx:\n",
    "                    label_ids.append(label[word_idx])\n",
    "                # For the other tokens in a word, we set the label to either the current label or -100, depending on\n",
    "                # the label_all_tokens flag.\n",
    "                else:\n",
    "                    label_ids.append(label[word_idx] if label_all_tokens else -100)\n",
    "                previous_word_idx = word_idx\n",
    "\n",
    "            labels.append(label_ids)\n",
    "\n",
    "        tokenized_inputs[\"labels\"] = labels\n",
    "        return tokenized_inputs\n",
    "    tokenized_datasets = raw_dataset.map(tokenize_and_align_labels, batched=True)\n",
    "    tokenized_datasets.set_format(type=\"torch\", columns=['attention_mask', 'input_ids', 'labels', 'token_type_ids'])\n",
    "\n",
    "    return tokenized_datasets, model, tokenizer\n",
    "\n",
    "\n",
    "# 训练\n",
    "def train_loop(seed, idx):\n",
    "    metric = load_metric(\"./char_ner_metric.py\")\n",
    "    tokenized_datasets, model, tokenizer = get_model(seed)\n",
    "    data_collator = DataCollatorForTokenClassification(tokenizer)\n",
    "\n",
    "    def compute_metrics(p):\n",
    "        predictions, labels = p\n",
    "        predictions = np.argmax(predictions, axis=2)\n",
    "\n",
    "        # Remove ignored index (special tokens)\n",
    "        true_predictions = [\n",
    "            [names[p] for (p, l) in zip(prediction, label) if l != -100]\n",
    "            for prediction, label in zip(predictions, labels)\n",
    "        ]\n",
    "        true_labels = [\n",
    "            [names[l] for (p, l) in zip(prediction, label) if l != -100]\n",
    "            for prediction, label in zip(predictions, labels)\n",
    "        ]\n",
    "\n",
    "        results = metric.compute(predictions=true_predictions, references=true_labels)\n",
    "        return {\n",
    "            \"precision\": results[\"overall_precision\"],\n",
    "            \"recall\": results[\"overall_recall\"],\n",
    "            \"f1\": results[\"overall_f1\"],\n",
    "            \"accuracy\": results[\"overall_accuracy\"],\n",
    "        }\n",
    "\n",
    "    args = TrainingArguments(\n",
    "        \"ner-3\",\n",
    "        do_train=True,\n",
    "        do_eval=True,\n",
    "        evaluation_strategy=\"epoch\",\n",
    "        per_device_train_batch_size=batch_size,\n",
    "        per_device_eval_batch_size=batch_size,\n",
    "        num_train_epochs=10,\n",
    "        gradient_accumulation_steps=gradient_accumulation_steps,\n",
    "        learning_rate=3e-5,\n",
    "        weight_decay=5e-3,\n",
    "        max_grad_norm=1.0,\n",
    "        warmup_ratio=0.2,\n",
    "        save_strategy=\"no\",\n",
    "        load_best_model_at_end=True,\n",
    "        no_cuda=False,\n",
    "        fp16=True,\n",
    "        label_smoothing_factor=0.1,\n",
    "        dataloader_pin_memory=True,\n",
    "    )\n",
    "\n",
    "    trainer = Trainer(\n",
    "        model,\n",
    "        args,\n",
    "        train_dataset=tokenized_datasets[\"train\"],\n",
    "        eval_dataset=tokenized_datasets[\"test\"],\n",
    "        tokenizer=tokenizer,\n",
    "        compute_metrics=compute_metrics\n",
    "    )\n",
    "\n",
    "    trainer.train()\n",
    "\n",
    "    model.save_pretrained(\"./ner_%d\"%idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d1d2a59b-ceda-4caa-afe9-8051dbe04b75",
   "metadata": {},
   "outputs": [
    {
     "ename": "ImportError",
     "evalue": "To be able to use this metric, you need to install the following dependencies['seqeval'] using 'pip install seqeval' for instance'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mImportError\u001b[0m                               Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-5-d4942db578a5>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mtrain_loop\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m402\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-4-e722ac3e53a0>\u001b[0m in \u001b[0;36mtrain_loop\u001b[0;34m(seed, idx)\u001b[0m\n\u001b[1;32m    120\u001b[0m \u001b[0;31m# 训练\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    121\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mtrain_loop\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mseed\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 122\u001b[0;31m     \u001b[0mmetric\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mload_metric\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"./seqeval.py\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    123\u001b[0m     \u001b[0mtokenized_datasets\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtokenizer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mseed\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    124\u001b[0m     \u001b[0mdata_collator\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mDataCollatorForTokenClassification\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtokenizer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/datasets/load.py\u001b[0m in \u001b[0;36mload_metric\u001b[0;34m(path, config_name, process_id, num_process, cache_dir, experiment_id, keep_in_memory, download_config, download_mode, script_version, **metric_init_kwargs)\u001b[0m\n\u001b[1;32m    613\u001b[0m         \u001b[0mdownload_config\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdownload_config\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    614\u001b[0m         \u001b[0mdownload_mode\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdownload_mode\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 615\u001b[0;31m         \u001b[0mdataset\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    616\u001b[0m     )\n\u001b[1;32m    617\u001b[0m     \u001b[0mmetric_cls\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mimport_main_class\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodule_path\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdataset\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/datasets/load.py\u001b[0m in \u001b[0;36mprepare_module\u001b[0;34m(path, script_version, download_config, download_mode, dataset, force_local_path, dynamic_modules_path, return_resolved_file_path, **download_kwargs)\u001b[0m\n\u001b[1;32m    453\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mneeds_to_be_installed\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    454\u001b[0m         raise ImportError(\n\u001b[0;32m--> 455\u001b[0;31m             \u001b[0;34mf\"To be able to use this {module_type}, you need to install the following dependencies\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    456\u001b[0m             \u001b[0;34mf\"{[lib_name for lib_name, lib_path in needs_to_be_installed]} using 'pip install \"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    457\u001b[0m             \u001b[0;34mf\"{' '.join([lib_path for lib_name, lib_path in needs_to_be_installed])}' for instance'\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mImportError\u001b[0m: To be able to use this metric, you need to install the following dependencies['seqeval'] using 'pip install seqeval' for instance'"
     ]
    }
   ],
   "source": [
    "train_loop(402, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44b5eab9-ce83-4f5a-ba32-fb12a7e428e2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
